import torch
from torch import nn
import torch.nn.functional as F
from math import sqrt
from transformers import AutoConfig
from transformers import AutoTokenizer


class AttentionHead(nn.Module):
    def __init__(self, embed_dim, head_dim):
        super().__init__()
        self.q = nn.Linear(embed_dim, head_dim)
        self.k = nn.Linear(embed_dim, head_dim)
        self.v = nn.Linear(embed_dim, head_dim)

    def forward(self, query, key, value, mask=None):
        query, key, value = self.q(query), self.k(key), self.v(value)
        scores = torch.bmm(query, key.transpose(1, 2)) / sqrt(query.size(-1))
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -float("inf"))
        weights = F.softmax(scores, dim=-1)
        return torch.bmm(weights, value)


class MultiHeadAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        embed_dim = config.hidden_size
        num_heads = config.num_attention_heads
        head_dim = embed_dim // num_heads
        self.heads = nn.ModuleList(
            [AttentionHead(embed_dim, head_dim) for _ in range(num_heads)]
        )
        self.output_linear = nn.Linear(embed_dim, embed_dim)

    def forward(self, query, key, value, mask=None, query_mask=None, key_mask=None):
        if query_mask is not None and key_mask is not None:
            mask = torch.bmm(query_mask.unsqueeze(-1), key_mask.unsqueeze(1))
        x = torch.cat([h(query, key, value, mask) for h in self.heads], dim=-1)
        x = self.output_linear(x)
        return x


class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.linear_1 = nn.Linear(config.hidden_size, config.intermediate_size)
        self.linear_2 = nn.Linear(config.intermediate_size, config.hidden_size)
        self.gelu = nn.GELU()
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, x):
        x = self.linear_1(x)
        x = self.gelu(x)
        x = self.linear_2(x)
        x = self.dropout(x)
        return x


class TransformerEncoderLayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.layer_norm_1 = nn.LayerNorm(config.hidden_size)
        self.layer_norm_2 = nn.LayerNorm(config.hidden_size)
        self.attention = MultiHeadAttention(config)
        self.feed_forward = FeedForward(config)

    def forward(self, x, mask=None):
        # Apply layer normalization and then copy input into query, key, value
        hidden_state = self.layer_norm_1(x)
        # Apply attention with a skip connection
        x = x + self.attention(hidden_state, hidden_state, hidden_state, mask=mask)
        # Apply feed-forward layer with a skip connection
        x = x + self.feed_forward(self.layer_norm_2(x))
        return x


class Embeddings(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.token_embeddings = nn.Embedding(config.vocab_size,
                                             config.hidden_size)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings,
                                                config.hidden_size)
        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=1e-12)
        self.dropout = nn.Dropout()

    def forward(self, input_ids):
        # Create position IDs for input sequence
        seq_length = input_ids.size(1)
        position_ids = torch.arange(seq_length, dtype=torch.long).unsqueeze(0)
        # Create token and position embeddings
        token_embeddings = self.token_embeddings(input_ids)
        position_embeddings = self.position_embeddings(position_ids)
        # Combine token and position embeddings
        embeddings = token_embeddings + position_embeddings
        embeddings = self.layer_norm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings


class TransformerEncoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.embeddings = Embeddings(config)
        self.layers = nn.ModuleList(
            [TransformerEncoderLayer(config) for _ in range(config.num_hidden_layers)]
        )

    def forward(self, x, mask=None):
        x = self.embeddings(x)
        for layer in self.layers:
            x = layer(x, mask)
        return x


model_ckpt = "../../dataroot/models/bert-base-uncased"


def f1():
    tokenizer = AutoTokenizer.from_pretrained(model_ckpt)

    text = "time flies like an arrow"
    inputs = tokenizer(text, return_tensors="pt", add_special_tokens=False)
    # print(inputs.input_ids)

    config = AutoConfig.from_pretrained(model_ckpt)
    token_emb = nn.Embedding(config.vocab_size, config.hidden_size)
    # Embedding(30522, 768)
    # BERT-base-uncased 模型对应的词表大小为 30522，每个词语的词向量维度为 768
    # print(token_emb)
    inputs_embeds = token_emb(inputs.input_ids)
    # 输入的文本time flies like an arrow，是1句话batch_size=1，是5个单词seq_len=5，模型的词向量维度为768即hidden_dim=768
    # torch.Size([1, 5, 768])
    # Embedding 层把输入的词语序列映射到了尺寸为 [batch_size, seq_len, hidden_dim] 的张量。
    '''
    tensor([[[-0.9596, -1.3469,  0.6265,  ..., -1.5737, -0.9345, -0.0862],
         [ 1.0920, -0.7868,  0.0566,  ...,  0.0331, -0.4390, -1.3692],
         [-1.2536, -0.8397, -0.2229,  ...,  0.1373, -1.3888,  3.1115],
         [-0.8423,  0.3149,  0.1182,  ...,  0.2165, -1.0138, -2.1425],
         [ 2.4727,  0.9180,  1.0240,  ..., -0.9168,  0.3953, -0.8255]]],
       grad_fn=<EmbeddingBackward0>)
    '''
    # print(inputs_embeds.size())
    # print(inputs_embeds)
    return inputs_embeds


def f2():
    tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
    config = AutoConfig.from_pretrained(model_ckpt)

    text = "time flies like an arrow"
    inputs = tokenizer(text, return_tensors="pt", add_special_tokens=False)

    encoder = TransformerEncoder(config)
    print(encoder(inputs.input_ids).size())


# scaled_dot_product_attention 简化版的 Scaled Dot-product Attention
def scaled_dot_product_attention(query, key, value, query_mask=None, key_mask=None, mask=None):
    dim_k = query.size(-1)
    scores = torch.bmm(query, key.transpose(1, 2)) / sqrt(dim_k)
    # torch.Size([1, 5, 5]),这里Q,K的序列长度都为 5，因此生成了一个5*5的注意力分数矩阵
    print(scores.size())
    '''
    上面的代码还考虑了Q，K，V序列的Mask。填充 (padding) 字符不应该参与计算，因此将对应的注意力分数设置为负无穷，
    这样 softmax 之后其对应的注意力权重就为 0 了（e的负无穷次幂=0）
    注意！上面的做法会带来一个问题：当Q，K序列相同时，注意力机制会为上下文中的相同单词分配非常大的分数（点积为 1），
    而在实践中，相关词往往比相同词更重要。例如对于上面的例子，只有关注“time”和“arrow”才能够确认“flies”的含义。
    因此，多头注意力 (Multi-head Attention) 出现了！
    '''
    if query_mask is not None and key_mask is not None:
        mask = torch.bmm(query_mask.unsqueeze(-1), key_mask.unsqueeze(1))
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -float("inf"))
    weights = F.softmax(scores, dim=-1)
    # tensor([[1., 1., 1., 1., 1.]], grad_fn=<SumBackward1>) 应用 Softmax 标准化注意力权重
    print(weights.sum(dim=-1))
    attn_outputs = torch.bmm(weights, value)
    # torch.Size([1, 5, 768]) 将注意力权重与value序列相乘
    print(attn_outputs.shape)
    return attn_outputs


def f3(inputs_embeds):
    q = k = v = inputs_embeds
    return scaled_dot_product_attention(q, k, v)

def f4():
    tokenizer = AutoTokenizer.from_pretrained(model_ckpt)

    text = "time flies like an arrow"
    inputs = tokenizer(text, return_tensors="pt", add_special_tokens=False)
    config = AutoConfig.from_pretrained(model_ckpt)
    token_emb = nn.Embedding(config.vocab_size, config.hidden_size)
    inputs_embeds = token_emb(inputs.input_ids)

    multihead_attn = MultiHeadAttention(config)
    query = key = value = inputs_embeds
    attn_output = multihead_attn(query, key, value)
    print(attn_output.size())

if __name__ == '__main__':
     f4()
